# -*- coding: utf-8 -*-
"""
   File Name：     ocr.py
   Description :  wrapping for face recognition class, instantiation occurs in modules/face_server/utils/api.py
   Author :       KangJiaHui
   date：         2021/02/04
"""
import itertools
import re
import os
import yaml
from datetime import datetime, timedelta
from paddleocr import PaddleOCR
from module.utils.utils import base64_to_image, image_to_base64
from module.utils.error import AimPose, NameRecFail, NameMatchFail
from module.face_recog.faceRecognition import FaceRecognition
import cv2
import requests

face = FaceRecognition()


def resize_img(image, _max=1500):
    if image.shape[0] > _max or image.shape[1] > _max:
        if image.shape[0] > image.shape[1]:
            proportion = _max / image.shape[0]
        else:
            proportion = _max / image.shape[1]
        image = cv2.resize(image, None, fx=proportion, fy=proportion, interpolation=cv2.INTER_AREA)
    return image


def if_exist(keyword_regular, result):
    """
    Match the regular expression of keyword and give all locations for each matched keyword.
    :param result: list, return of get_all().
    :param keyword_regular: string, could be regular expression. e.x."姓名", "危险｜爆炸"
    :return: list, e.x.[[[[1299.0, 480.0], [1448.0, 477.0], [1449.0, 514.0], [1300.0, 517.0]], ('危险', 0.99250406)],
                    [[[1663.0, 480.0], [1886.0, 480.0], [1886.0, 510.0], [1663.0, 510.0]], ('爆炸', 0.8288886)]]
    """
    for line in result:
        if re.search(keyword_regular, line[-1][0]):
            return True
    return False


class OCR(object):
    def __init__(self):
        self.model = PaddleOCR(use_angle_cls=True, lang="ch")
        # self.my_model = PaddleOCR(det_model_dir='module/params/my_ch_db_res18',
        #                           rec_model_dir='module/params/ch_ppocr_mobile_v2.0_rec_infer',
        #                           rec_char_dict_path='module/params/ppocr_keys_v1.txt',
        #                           cls_model_dir='module/params/ch_ppocr_mobile_v2.0_cls_infer',
        #                           use_angle_cls=True)
        with open(os.path.join(os.getcwd(), "confs/certificate.yaml"), "r", encoding='UTF-8') as f:
            self.config = yaml.load(f, Loader=yaml.FullLoader)

    def get_all(self, img):
        """
        Recognize all information from an image.
        :param img: numpy.ndarray
        :return: list, e.x.[[[[1319.0, 171.0], [1673.0, 173.0], [1673.0, 232.0], [1319.0, 230.0]], ('服务单位', 0.9971724)]
        , [[[333.0, 201.0], [414.0, 201.0], [414.0, 248.0], [333.0, 248.0]], ('姓名', 0.9986272)]]
        """
        result = self.model.ocr(img, cls=True)
        # print(result)
        return result

    def id_recog(self, result):
        """
        Recognize witch certificate it is, according to the OCR result of image.
        :param result: list, the OCR result of input image.
        :return: string, the key of matched certificate.
        """
        for key in self.config:
            regular = self.config[key]["关键字"]
            flag = True
            for i in regular:
                if not if_exist(i, result):
                    flag = False
            if flag:
                return key
        return None

    def recog_one(self, image_base64):
        """
        Recognize one image, and to update tmp_result dict
        :param paper_name: str, the name of certificate. e.x. "道路危险货物运输驾驶员证"
        :param face_flag: bool, if True, extract face feature from input image.
        :param image_base64: image encoded in base64
        :return: result: e.x. {"证件名": "中华人民共和国道路运输证", "有效期": None, "经营范围": "危险化学品", "车牌": "鲁FBR932"}
        """
        image = base64_to_image(image_base64)
        image = resize_img(image)
        ocr_result = self.get_all(image)
        # print(ocr_result)
        id = self.id_recog(ocr_result)
        if not id:
            raise NameRecFail
        result = {}
        for key in self.config[id]:
            if key == "关键字" or key == "face_flag":
                continue
            result[key] = self.config[id][key]
        if self.config[id]["face_flag"]:
            feature, bbox, face_img = face.face_register(image_base64, 0.5)
            result["人脸"] = feature
            result["人脸位置"] = bbox  # [left, top, right, bottom]
            result["切图"] = face_img
        return result

    def face_match(self, image_base64, face_feature):
        """
        Match face in input image with the faces in all certificates. Once a certificate matched, return True.
        :param face_feature:
        :param image_base64: image encoded in base64
        :return: bool, True means face matched, False means face not matched.
        """
        return face.match_identity(face_feature, 0.5, 0.5, image_base64)

    @staticmethod
    def face_register(image_base64):
        """
        Extract face feature vector of input image.
        :param image_base64: image encoded in base64
        :return: list, face feature vector of input image.
        """
        return face.face_register(image_base64, 0.5)
